#!/bin/bash
gen_image_repo_url() {
    local IMAGE_NAME="$1"
    local LMD_IMAGE_NAME="$2"
    local DOCKER_APP_VERSION="$3"
    local ARCH_AND_GPU_INFO="cpu"
    
    # 1.0.0-gpu-cu121
    if ! command -v nvidia-smi &> /dev/null; then
        echo "No NVIDIA driver detected."
        if command -v rocm-smi &> /dev/null && rocm-smi | grep -i "amd" > /dev/null; then
            echo "Found AMD GPU"
            ARCH_AND_GPU_INFO="rocm"
        else
            echo "only cpu"
            ARCH_AND_GPU_INFO="cpu"
            # ARCH_AND_GPU_INFO="mac-arm64"
        fi
    else
        if ! command -v nvcc &> /dev/null; then
            echo "Error: The nvcc command was not found. Make sure you have the CUDA Toolkit installed."
        else
            echo "Found NVIDIA drivers and cuda"
            nvcc_output=$(nvcc --version)
            if [[ $? -ne 0 ]]; then
                echo "Error: Failed to execute 'nvcc --version'." >&2
                exit 1
            fi

            cuda_line=$(echo "$nvcc_output" | grep 'Cuda compilation tools')

            if [[ $cuda_line =~ release\ ([0-9]+\.[0-9]+) ]]; then
                cuda_version=${BASH_REMATCH[1]}
                echo "CUDA Version: $cuda_version"
            else
                echo "Error: Unable to parse CUDA version from output." >&2
                exit 1
            fi

            if [ -z "$cuda_version" ]; then
                echo "Error: Unable to extract CUDA version number from nvcc output."
                ARCH_AND_GPU_INFO="gpu-cu126"
            else
                if [[ "$cuda_version" == "12.8" || "$cuda_version" == "12.9" ]]; then
                    ARCH_AND_GPU_INFO="gpu-cu128"
                elif [[ "$cuda_version" == "12.6" || "$cuda_version" == "12.7" ]]; then
                    ARCH_AND_GPU_INFO="gpu-cu126"
                elif [[ "$cuda_version" == "12.4" || "$cuda_version" == "12.5" ]]; then
                    ARCH_AND_GPU_INFO="gpu-cu124"
                elif [[ "$cuda_version" == "12.0" || "$cuda_version" == "12.1" || "$cuda_version" == "12.2" || "$cuda_version" == "12.3" ]]; then
                    ARCH_AND_GPU_INFO="gpu-cu121"
                else
                    ARCH_AND_GPU_INFO="gpu-cu121"
                fi
            fi
        fi
    fi
    
    local LMD_REGISTRY="crpi-ru0fy7yg2nq6feyc.cn-shanghai.personal.cr.aliyuncs.com/3rd-party"
    IMAGE_TAG="${DOCKER_APP_VERSION}-${ARCH_AND_GPU_INFO}"
    export APP_IMAGE_REPO_RESULT="$LMD_REGISTRY/$LMD_IMAGE_NAME:$IMAGE_TAG"
    
}
